package mxp

import cpu74.ASM74Lexer
import cpu74.ASM74Parser
import org.antlr.v4.runtime.ANTLRInputStream
import org.antlr.v4.runtime.CommonTokenStream
import org.antlr.v4.runtime.ParserRuleContext
import java.io.File
import java.io.FileOutputStream
import java.io.OutputStream
import java.nio.file.Files
import java.nio.file.Paths
import java.util.*
import kotlin.system.exitProcess

private val model = MCodeDefModel.fromJson(Files.readAllBytes(Paths.get("../mcode.json")).toString(charset("UTF-8")))

private val romAddrMap = mutableMapOf<String, Int>()
private val romAddrReverseMap = mutableMapOf<Int, MutableList<String>>()
private val romAddrHolderMap = mutableMapOf<Int, String>()
private val ramAddrMap = mutableMapOf<String, Int>()
private val rom = Array(655350) { 0 }

private var rom_rodata = Pair(0, 0)
private var rom_data = Pair(0, 0)
private var rom_code = Pair(0, 0)
private var ram_bss = Pair(0, 0)
private var ram_data = Pair(0, 0)

private var romPos = 0
private var ramPos = 0xA000

private var fos: OutputStream? = null

private val exprCtx = object : HashMap<String, Int>() {
    override fun get(key: String): Int? {
        if (romAddrMap.containsKey(key)) {
            return romAddrMap[key]
        } else if (ramAddrMap.containsKey(key)) {
            return ramAddrMap[key]
        } else {
            return null
        }
    }
}

private fun evalExpr2(expr2Ctx: ASM74Parser.Expr2Context): Int {
    if (expr2Ctx.NUMBER() != null) {
        return str2int(expr2Ctx.NUMBER().text)
    } else if (expr2Ctx.ID() != null) {
        if (expr2Ctx.expr() != null) {
            when (expr2Ctx.ID().text) {
                "H" -> {
                    return (evalExpr(expr2Ctx.expr()) shr 8) and 0xff
                }
                "L" -> {
                    return evalExpr(expr2Ctx.expr()) and 0xff
                }
            }
        } else {
            return exprCtx[expr2Ctx.ID().text]!!
        }
    }

    println("Impossible")
    exitProcess(1)
}

private fun evalExpr1(expr1Ctx: ASM74Parser.Expr1Context): Int {
    var result = 0
    var op = ""
    for (child in expr1Ctx.children) {
        if (child.javaClass == ASM74Parser.Expr2Context::class.java) {
            if (op == "") {
                result = evalExpr2(child as ASM74Parser.Expr2Context)
            } else {
                if (op == "*") {
                    result *= evalExpr2(child as ASM74Parser.Expr2Context)
                } else if (op == "/") {
                    result /= evalExpr2(child as ASM74Parser.Expr2Context)
                }
            }
        } else {
            op = child.text
        }
    }

    return result
}

private fun evalExpr(exprCtx: ASM74Parser.ExprContext): Int {
    var result = 0
    var op = ""
    for (child in exprCtx.children) {
        if (child.javaClass == ASM74Parser.Expr1Context::class.java) {
            if (op == "") {
                result = evalExpr1(child as ASM74Parser.Expr1Context)
            } else {
                if (op == "+") {
                    result += evalExpr1(child as ASM74Parser.Expr1Context)
                } else if (op == "-") {
                    result -= evalExpr1(child as ASM74Parser.Expr1Context)
                }
            }
        } else {
            op = child.text
        }
    }

    return result
}

private fun calcAddr_rodata(progCtxs: List<ASM74Parser.ProgContext>) {
    for (progCtx in progCtxs) {
        var isReadOnlyDataSegment = false
        for (lineCtx in progCtx.line()) {
            if (isReadOnlyDataSegment && lineCtx.label() != null) {
                romAddrMap[lineCtx.label().ID().text] = romPos
            }

            if (lineCtx.cmd() != null) {
                val cmdCtx = lineCtx.cmd()
                val cmd = cmdCtx.start.text
                if (cmd == ".segment") {
                    isReadOnlyDataSegment = lineCtx.cmd().STRING().text.trim('"').toUpperCase() == "RODATA"
                }

                if (isReadOnlyDataSegment) {
                    if (cmd == ".byte") {
                        val bytes = cmdCtx.NUMBER().map { str2int(it.text) }
                        bytes.forEach { rom[romPos++] = (0xff and it) }
                    } else if (cmd == "=") {
                        val addr = evalExpr(cmdCtx.expr())
                        rom[romPos++] = addr and 0xff
                        rom[romPos++] = (addr shr 8) and 0xff
                    }
                }
            }
        }
    }
}

private fun calcAddr_bss(progCtxs: List<ASM74Parser.ProgContext>) {
    for (progCtx in progCtxs) {
        var isBssSegment = false
        for (lineCtx in progCtx.line()) {
            if (isBssSegment && lineCtx.label() != null) {
                ramAddrMap[lineCtx.label().ID().text] = ramPos
            }

            if (lineCtx.cmd() != null) {
                val cmdCtx = lineCtx.cmd()
                val cmd = cmdCtx.start.text
                if (cmd == ".segment") {
                    isBssSegment = lineCtx.cmd().STRING().text.trim('"').toUpperCase() == "BSS"
                }

                if (isBssSegment) {
                    if (cmd == ".res") {
                        ramPos += str2int(lineCtx.cmd().NUMBER(0).text)
                    }
                }
            }
        }
    }
}

private fun calcAddr_data(progCtxs: List<ASM74Parser.ProgContext>) {
    for (progCtx in progCtxs) {
        var isDataSegment = false
        for (lineCtx in progCtx.line()) {
            if (isDataSegment && lineCtx.label() != null) {
                val name = lineCtx.label().ID().text
                romAddrMap["_init_$name"] = romPos
                ramAddrMap[name] = ramPos
            }

            if (lineCtx.cmd() != null) {
                val cmdCtx = lineCtx.cmd()
                val cmd = cmdCtx.start.text
                if (cmd == ".segment") {
                    isDataSegment = lineCtx.cmd().STRING().text.trim('"').toUpperCase() == "DATA"
                }

                if (isDataSegment) {
                    when (cmd) {
                        ".byte" -> {
                            val bytes = cmdCtx.NUMBER().map { str2int(it.text) }
                            bytes.forEach {
                                rom[romPos++] = 0xff and it
                            }
                            ramPos += bytes.size
                        }
                        ".word" -> {
                            val words = cmdCtx.NUMBER().map { str2int(it.text) }
                            words.forEach {
                                rom[romPos++] = 0xff and it
                                rom[romPos++] = 0xff and (it shr 8)
                            }
                            ramPos += words.size * 2
                        }
                        ".dword" -> {
                            val dwords = cmdCtx.NUMBER().map { str2int(it.text) }
                            dwords.forEach {
                                rom[romPos++] = 0xff and it
                                rom[romPos++] = 0xff and (it shr 8)
                                rom[romPos++] = 0xff and (it shr 16)
                                rom[romPos++] = 0xff and (it shr 24)
                            }
                            ramPos += dwords.size * 4
                        }
                        ".addr" -> {
                            romAddrHolderMap[romPos] = cmdCtx.ID().text
                            ramPos += 2
                            romPos += 2
                        }
                        ".res" -> {
                            val size = str2int(lineCtx.cmd().NUMBER(0).text)
                            ramPos += size
                        }
                    }
                }
            }
        }
    }
}

private val AMCtx2NameMap = mapOf(
    ASM74Parser.Am_defaultContext::class.java to "default",
    ASM74Parser.Am_immContext::class.java to "imm",
    ASM74Parser.Am_directContext::class.java to "direct",
    ASM74Parser.Am_indirectContext::class.java to "indirect",
    ASM74Parser.Am_direct_indexContext::class.java to "direct_index",
    ASM74Parser.Am_indirect_indexContext::class.java to "indirect_index",
    ASM74Parser.Am_index_indirectContext::class.java to "index_indirect"
)

private fun getAmName(amCtx: ASM74Parser.AmContext): String? =
    AMCtx2NameMap[(amCtx.children[0] as ParserRuleContext).javaClass]

private fun getLabelsAt(addr: Int): List<String> = romAddrReverseMap.getOrDefault(addr, emptyList())

private fun calcAddr_code(progCtxs: List<ASM74Parser.ProgContext>) {
    for (progCtx in progCtxs) {
        var isCodeSegment = false
        for (lineCtx in progCtx.line()) {
            if (isCodeSegment && lineCtx.label() != null) {
                romAddrMap[lineCtx.label().ID().text] = romPos
            }

            if (lineCtx.cmd() != null) {
                val cmdCtx = lineCtx.cmd()
                val cmd = cmdCtx.start.text
                if (cmd == ".segment") {
                    isCodeSegment = lineCtx.cmd().STRING().text.trim('"').toUpperCase() == "CODE"
                }

                if (isCodeSegment && cmd == ".proc") {
                    romAddrMap[cmdCtx.ID().text] = romPos
                }
            }

            if (isCodeSegment && lineCtx.opcode() != null) {
                var amName = getAmName(lineCtx.am())
                if (amName == "direct_index") amName = "direct"
                else if (amName == "index_indirect") amName = "indirect"

                romPos += model.getInstructionLength(lineCtx.opcode().text, amName)
            }
        }
    }
}

private fun calcAddr(progCtxs: List<ASM74Parser.ProgContext>) {
    val codeSegmentBegin = romPos
    calcAddr_code(progCtxs)
    val codeSegmentEnd = romPos

    romPos += 0x10 - (romPos and 0xf)

    val rodataSegmentBegin = romPos
    calcAddr_rodata(progCtxs)
    val rodataSegmentEnd = romPos

    romPos += 0x10 - (romPos and 0xf)

    val ramBssBegin = ramPos
    calcAddr_bss(progCtxs)
    val ramBssEnd = ramPos

    val dataSegmentBegin = romPos
    val ramDataBegin = ramPos
    calcAddr_data(progCtxs)
    val ramDataEnd = ramPos
    val dataSegmentEnd = romPos

    rom_rodata = Pair(rodataSegmentBegin, rodataSegmentEnd)
    rom_data = Pair(dataSegmentBegin, dataSegmentEnd)
    rom_code = Pair(codeSegmentBegin, codeSegmentEnd)

    ram_bss = Pair(ramBssBegin, ramBssEnd)
    ram_data = Pair(ramDataBegin, ramDataEnd)

    romAddrMap.forEach { label, addr ->
        romAddrReverseMap.getOrPut(addr, { mutableListOf() }).add(label)
    }

    romAddrHolderMap.forEach { addr, label ->
        rom[addr] = 0xff and romAddrMap[label]!!
        rom[addr + 1] = 0xff and (romAddrMap[label]!! shr 8)
    }
}

private fun asm_proc_line(lineCtx: ASM74Parser.LineContext, startPos: Int): Int {
    val opcode = lineCtx.opcode().text
    val amCtx = lineCtx.am()
    val amName = getAmName(lineCtx.am())
    val bytes = when (amName) {
        "default" ->
            model.getBinaryCode(opcode, "default", 0, 0)
        "imm" ->
            model.getBinaryCode(opcode, "imm", 0, evalExpr(amCtx.am_imm().num().expr()))
        "direct" ->
            model.getBinaryCode(opcode, "direct", evalExpr(amCtx.am_direct().address().expr()), 0)
        "indirect" ->
            model.getBinaryCode(opcode, "indirect", evalExpr(amCtx.am_indirect().address().expr()), 0)
        "indirect_index" -> {
            val value = evalExpr(amCtx.am_indirect_index().num().expr())
            // TODO 这里导致代码段长度不好计算，暂时先注释掉
//            if (value == 0) {
//                model.getBinaryCode(
//                    opcode,
//                    "indirect",
//                    evalExpr(amCtx.am_indirect_index().address().expr()),
//                    0
//                )
//            } else {
            model.getBinaryCode(
                opcode,
                "indirect_index",
                evalExpr(amCtx.am_indirect_index().address().expr()),
                value
            )
//            }

        }
        "direct_index" ->
            model.getBinaryCode(
                opcode,
                "direct",
                evalExpr(amCtx.am_direct_index().address().expr()) + evalExpr(amCtx.am_direct_index().num().expr()),
                0
            )
        "index_indirect" ->
            model.getBinaryCode(
                opcode,
                "indirect",
                evalExpr(amCtx.am_index_indirect().address().expr()) + evalExpr(amCtx.am_index_indirect().num().expr()),
                0
            )
        else -> {
            println("unknown addressing mode")
            exitProcess(1)
        }
    }

    var s = String.format("%04x\t", startPos)
    for (i in 0 until bytes.size) {
        s += String.format("%02x ", bytes[i])
    }
    s = s.padEnd(20, ' ')
    s += getLabelsAt(startPos).joinToString(",")
    s = s.padEnd(40, ' ')
    s += " " + opcode + " " + amCtx.text
    output_list(s)

    var pos = startPos
    bytes.forEach { rom[pos++] = 0xff and it.toInt() }

    return bytes.size
}

private fun asm(progCtxs: List<ASM74Parser.ProgContext>) {
    for (progCtx in progCtxs) {
        var isCodeSegment = false
        var pos = -1
        for (lineCtx in progCtx.line()) {
            if (lineCtx.cmd() != null) {
                val cmdCtx = lineCtx.cmd()
                val cmd = cmdCtx.start.text
                if (cmd == ".segment") {
                    isCodeSegment = lineCtx.cmd().STRING().text.trim('"').toUpperCase() == "CODE"
                }

                if (isCodeSegment && cmd == ".proc") {
                    pos = romAddrMap[cmdCtx.ID().text]!!
                }
            }

            if (isCodeSegment) {
                if (lineCtx.opcode() != null) {
                    pos += asm_proc_line(lineCtx, pos)
                }
            }
        }
    }
}

private fun genCode_init_data(progCtxs: List<ASM74Parser.ProgContext>): String {
    var label = 0
    var init_data = ".segment \"CODE\"\n"
    init_data += ".proc _init_data\n"

    for (progCtx in progCtxs) {
        val varSizeMap = mutableMapOf<String, Int>()
        var isDataSegment = false
        var name = ""
        for (lineCtx in progCtx.line()) {
            if (isDataSegment && lineCtx.label() != null) {
                name = lineCtx.label().ID().text
                varSizeMap.put(name, 0)
            }

            if (lineCtx.cmd() != null) {
                val cmdCtx = lineCtx.cmd()
                val cmd = cmdCtx.start.text
                if (cmd == ".segment") {
                    isDataSegment = lineCtx.cmd().STRING().text.trim('"').toUpperCase() == "DATA"
                }

                if (isDataSegment) {
                    when (cmd) {
                        ".byte" -> {
                            varSizeMap[name] = varSizeMap[name]?.plus(cmdCtx.NUMBER().size)!!
                        }
                        ".word" -> {
                            varSizeMap[name] = varSizeMap[name]?.plus(cmdCtx.NUMBER().size * 2)!!
                        }
                        ".dword" -> {
                            varSizeMap[name] = varSizeMap[name]?.plus(cmdCtx.NUMBER().size * 4)!!
                        }
                        ".addr" -> {
                            varSizeMap[name] = varSizeMap[name]?.plus(2)!!
                        }
                    }
                }
            }
        }

        varSizeMap.forEach { varname, size ->
            if (size <= 16) {
                for (i in 0 until size) {
                    init_data += """
                        lda _init_${varname},#$i
                        sta $varname,#$i
                        
                    """.trimIndent()
                }
            } else {
                val label0 = "_init_data_${label++}"
                val label1 = "_init_data_${label++}"
                init_data += """
                    lda #${size and 0xff}
                    sta TMP0
                    lda #${(size shr 8) and 0xff}
                    sta TMP0,#1
                    lda #L(_init_$varname)
                    sta TMP1
                    lda #H(_init_$varname)
                    sta TMP1,#1
                    lda #L($varname)
                    sta TMP2
                    lda #H($varname)
                    sta TMP2,#1
           $label0: lda TMP0
                    or TMP0,#1
                    jz $label1
                    lda (TMP1)
                    sta (TMP2)
                    sec
                    lda TMP0
                    sub #1
                    sta TMP0
                    lda TMP0,#1
                    sub #0
                    sta TMP0,#1
                    clc
                    lda TMP1
                    add #1
                    sta TMP1
                    lda TMP1,#1
                    add #0
                    sta TMP1,#1
                    clc
                    lda TMP2
                    add #1
                    sta TMP2
                    lda TMP2,#1
                    add #0
                    sta TMP2,#1
                    jmp $label0
           $label1:

                """.trimIndent()
            }
        }
    }
    init_data += "jmp _main\n"

    return init_data
}

private fun genCode_entry(): String {
    return """
        .segment "CODE"
        .proc _init
        lda #L(STACK_BASE)
        sta SP
        lda #H(STACK_BASE)
        sta SP,#1
        jmp _init_data
        
    """.trimIndent()
}

private fun antlr_parse(content: String): ASM74Parser.ProgContext {
    val lexer = ASM74Lexer(ANTLRInputStream(content))
    val token = CommonTokenStream(lexer)
    val parser = ASM74Parser(token)
    return parser.prog()
}

private fun dumpRom(range: Pair<Int, Int>) {
    var n = 0
    val bytes = ByteArray(16)
    for (addr in range.first until range.second) {
        if (n == 0) {
            output_list("%04x\t".format(addr), false)
        }
        bytes[n % 16] = if (rom[addr] < 0x20) '.'.toByte() else rom[addr].toByte()
        output_list("%02x ".format(rom[addr]), false)

        n++
        if (n == 16) {
            n = 0
            output_list("\t${bytes.toString(charset("UTF-8"))}")
        }
    }
    if (n != 0) {
        output_list("${"".padEnd(3 * (16 - n))}\t${bytes.copyOf(n).toString(charset("UTF-8"))}")
    }
}


private fun output_list(str: String, cr: Boolean = true) {
    fos?.write(str.toByteArray(charset("UTF-8")))
    if (cr) {
        fos?.write('\n'.toInt())
    }

    print(str)
    if (cr) {
        println()
    }
}

private fun main(args: Array<String>) {
//    val afiles = arrayOf("../hello.s")
    val afiles = arrayOf("../cc74/test/table.s")
//    val afiles = arrayOf(
//        "../tinybasic/str.s",
//        "../tinybasic/io.s",
//        "../tinybasic/code.s",
//        "../tinybasic/lexer.s",
//        "../tinybasic/interpreter.s",
//        "../tinybasic/mxpbasic.s"
//    )
    val files = afiles.map { File(it).absolutePath }.sorted()


    val dir = File(files[0]).parent
    val dst = if (files.size == 1) files[0].replaceAfterLast('.', "bin", "$dir/out.bin") else "$dir/out.bin"
    val lst = if (files.size == 1) files[0].replaceAfterLast('.', "alst", "$dir/out.alst") else "$dir/out.alst"

    fos = FileOutputStream(lst)

    val progCtxs = mutableListOf<ASM74Parser.ProgContext>()
    for (file in files) {
        val content = Files.readAllBytes(Paths.get(file)).toString(charset("UTF-8"))
        progCtxs.add(antlr_parse(content))
    }

    progCtxs.add(0, antlr_parse(genCode_entry()))
    progCtxs.add(1, antlr_parse(genCode_init_data(progCtxs)))
    calcAddr(progCtxs)

    ramAddrMap["A"] = 0x8000
    ramAddrMap["B"] = 0x8004
    ramAddrMap["SP"] = 0x8008
    ramAddrMap["TMP0"] = 0x8010
    ramAddrMap["TMP1"] = 0x8020
    ramAddrMap["TMP2"] = 0x8030
    ramAddrMap["TMP3"] = 0x8040
    ramAddrMap["TMP4"] = 0x8050
    ramAddrMap["STACK_BASE"] = 0x9FFF

    asm(progCtxs)

    output_list("    \t==================== RODATA ===================")
    output_list("    \t 0  1  2  3  4  5  6  7  8  9  A  B  C  D  E  F")
    dumpRom(rom_rodata)
    output_list("    \t 0  1  2  3  4  5  6  7  8  9  A  B  C  D  E  F")
    output_list("    \t====================  DATA  ===================")
    output_list("    \t 0  1  2  3  4  5  6  7  8  9  A  B  C  D  E  F")
    dumpRom(rom_data)
    output_list("    \t 0  1  2  3  4  5  6  7  8  9  A  B  C  D  E  F")

    val segmentInfo = """
        ROM RODATA: [0x${rom_rodata.first.toString(16)}, 0x${rom_rodata.second.toString(16)}) ${rom_rodata.second - rom_rodata.first}
        ROM   DATA: [0x${rom_data.first.toString(16)}, 0x${rom_data.second.toString(16)}) ${rom_data.second - rom_data.first}
        ROM   CODE: [0x${rom_code.first.toString(16)}, 0x${rom_code.second.toString(16)}) ${rom_code.second - rom_code.first}

        RAM    BSS: [0x${ram_bss.first.toString(16)}, 0x${ram_bss.second.toString(16)}) ${ram_bss.second - ram_bss.first}
        RAM   DATA: [0x${ram_data.first.toString(16)}, 0x${ram_data.second.toString(16)}) ${ram_data.second - ram_data.first}

    """.trimIndent()
    println(segmentInfo)

    val output = rom2logisim(rom)
    Files.write(Paths.get(dst), output.toByteArray(charset("UTF-8")))

    fos?.close()
}